import torch
from torch import nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import  train_test_split


data = pd.read_csv("D:\learn\深度学习\day53_pytorch入门(二)\代码\dataset\HR.csv")

salaryMap = {"low": 1, "medium": 2, "high": 3}
data.salary = data.salary.map(salaryMap)
data = data.join(pd.get_dummies(data.part))
Y_data = data.left.values.reshape(-1, 1)
Y = torch.from_numpy(Y_data).type(torch.FloatTensor)
data.drop(columns=['part', 'left'], inplace=True)
X_data = data.replace(False, 0).replace(True, 1).values
X = torch.from_numpy(X_data).type(torch.float32)
print(X.shape)


class HRModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin_1 = nn.Linear(18, 128)
        self.lin_2 = nn.Linear(128, 64)
        self.lin_3 = nn.Linear(64, 1)
        self.activate = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, input):
        # 定义前向传播
        t = self.lin_1(input)
        t = self.activate(t)
        t = self.lin_2(t)
        t = self.activate(t)
        t = self.lin_3(t)
        t = self.sigmoid(t)
        return t


lr = 0.001


def get_model():
    model1 = HRModel()
    return model1, torch.optim.Adam(model1.parameters(), lr=lr)


loss_fn = nn.BCELoss()
model, opt = get_model()
batch_size = 64
steps = len(data) // batch_size
epochs = 500
r"""
for epoch in range(epochs):
    for i in range(steps):
        start = i * batch_size
        end = start + batch_size
        x = X[start:end]
        y = Y[start: end]
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    print('epoch:', epoch, ' ', 'loss:', loss_fn(model(X), Y))

res = ((model(X).data.numpy() > 0.5) == Y.numpy()).mean()
print(res)


HR_ds = TensorDataset(X, Y)
HR_dl = DataLoader(HR_ds, batch_size=batch_size)
for epoch in range(epochs):
    for x, y in HR_dl:
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    print('epoch:', epoch, ' ', 'loss:', loss_fn(model(X), Y))

train_x, test_x, train_y, test_y = train_test_split(X_data, Y_data, random_state=5)

train_x = torch.from_numpy(train_x).type(torch.FloatTensor)
test_x = torch.from_numpy(test_x).type(torch.FloatTensor)

train_y = torch.from_numpy(train_y).type(torch.FloatTensor)
test_y = torch.from_numpy(test_y).type(torch.FloatTensor)

train_dl = DataLoader(TensorDataset(train_x, train_y), batch_size= batch_size, shuffle=True)
test_dl = DataLoader(TensorDataset(test_x, test_y), batch_size=batch_size * 2, shuffle=True)




for epoch in range(epochs):
    model.train()
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_fn(pred, yb)
        loss.backward()
        opt.zero_grad()
        opt.step()

    model.eval()
    with torch.no_grad():
        acc_mean = np.mean([accuracy(model(xt), yt) for xt, yt in test_dl])
        print(acc_mean)
"""

def accuracy(out, yb1):
    return ((out.data.numpy() > 0.6) == yb1.numpy()).mean()


def loss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)

    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()
    return loss.item(), len(xb)


def fit(epochs, model, loss_fn, opt, train_dl, valid_dl):
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            loss_batch(model, loss_fn, xb, yb, opt)

        model.eval()
        with torch.no_grad():
            loss_sum = 0
            num_sum = 0
            for xb, yb in valid_dl:
                losses, nums = loss_batch(model, loss_fn, xb, yb)
                loss_sum += losses
                num_sum += nums
            acc_mean = np.mean([accuracy(model(xt), yt) for xt, yt in valid_dl])
            print(epoch, loss_sum / num_sum, acc_mean)



def get_data(train_ds, valid_ds, batch_size):
    return (
            DataLoader(train_ds, batch_size=batch_size, shuffle=True),
            DataLoader(valid_ds, batch_size=batch_size*2)
            )

train_x, test_x, train_y, test_y = train_test_split(X_data, Y_data, random_state=5)

train_x = torch.from_numpy(train_x).type(torch.FloatTensor)
test_x = torch.from_numpy(test_x).type(torch.FloatTensor)

train_y = torch.from_numpy(train_y).type(torch.FloatTensor)
test_y = torch.from_numpy(test_y).type(torch.FloatTensor)

train_dl, valid_dl = get_data(TensorDataset(train_x, train_y),TensorDataset(test_x, test_y), batch_size)
model, opt = get_model()
fit(epochs, model, loss_fn, opt, train_dl, valid_dl)



